import os
import shutil
import pydicom
from collections import defaultdict

def split_series_time(input_dir, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 遍历所有子文件夹
    for root, dirs, files in os.walk(input_dir):
        # 只处理包含dcm的文件夹
        dcm_files = [f for f in files if f.lower().endswith(".dcm")]
        if not dcm_files:
            continue

        folder_name = os.path.basename(root)
        print(f"处理文件夹: {folder_name}")

        # 根据SeriesInstanceUID分组
        series_dict = defaultdict(list)
        for f in dcm_files:
            f_path = os.path.join(root, f)
            try:
                ds = pydicom.dcmread(f_path, stop_before_pixels=True)
                series_uid = ds.SeriesInstanceUID
                series_dict[series_uid].append(f_path)
            except Exception as e:
                print(f"读取失败 {f_path}: {e}")


        if len(series_dict) <= 1:
            print('只有一个序列, 跳过')
            continue

        series_times = {}
        for uid, files in series_dict.items():
            times = []
            for f in files:
                try:
                    ds = pydicom.dcmread(f, stop_before_pixels=True, force=True)
                    if hasattr(ds, 'AcquisitionTime'):
                        times.append(ds.AcquisitionTime)
                except:
                    pass

            series_times[uid] = min(times) if times else "0"

        latest_uid = min(series_times, key=lambda k: series_times[k])

        # 保存到输出文件夹
        for idx, (series_uid, files_in_series) in enumerate(series_dict.items(), start=1):
            new_folder_name = f"{folder_name}_{idx}"
            if series_uid == latest_uid:
                new_folder_name += '_earliest'
            new_folder_path = os.path.join(output_dir, new_folder_name)
            os.makedirs(new_folder_path, exist_ok=True)

            for f in files_in_series:
                shutil.move(f, new_folder_path)

            print(f"  -> 保存到 {new_folder_name}, 共 {len(files_in_series)} 个文件")


def split_series_thickness(input_dir, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 遍历所有子文件夹
    for root, dirs, files in os.walk(input_dir):
        # 只处理包含dcm的文件夹
        dcm_files = [f for f in files if f.lower().endswith(".dcm")]
        if not dcm_files:
            continue

        folder_name = os.path.basename(root)
        print(f"处理文件夹: {folder_name}")

        # 根据SeriesInstanceUID分组
        series_dict = defaultdict(list)
        for f in dcm_files:
            f_path = os.path.join(root, f)
            try:
                ds = pydicom.dcmread(f_path, stop_before_pixels=True)
                series_uid = ds.SeriesInstanceUID
                series_dict[series_uid].append(f_path)
            except Exception as e:
                print(f"读取失败 {f_path}: {e}")


        if len(series_dict) <= 1:
            print('只有一个序列, 跳过')
            continue

        series_thickness = {}
        for uid, files in series_dict.items():
            thickness = []
            for f in files:
                try:
                    ds = pydicom.dcmread(f, stop_before_pixels=True, force=True)
                    if hasattr(ds, 'SliceThickness'):
                        thickness.append(ds.SliceThickness)
                except:
                    pass

            series_thickness[uid] = thickness[0] if thickness else "0"

        # 保存到输出文件夹
        for idx, (series_uid, files_in_series) in enumerate(series_dict.items(), start=1):
            new_folder_name = f"{folder_name}_{idx}_{series_thickness[series_uid]}mm"

            new_folder_path = os.path.join(output_dir, new_folder_name)
            os.makedirs(new_folder_path, exist_ok=True)

            for f in files_in_series:
                shutil.copy2(f, new_folder_path)

            print(f"  -> 保存到 {new_folder_name}, 共 {len(files_in_series)} 个文件")


if __name__ == "__main__":
    input_dir = r"/media/imed/d9f520b2-6e7a-4842-af77-fba3ad2d8d27/test"   # 顶层文件夹
    output_dir = r"/media/imed/d9f520b2-6e7a-4842-af77-fba3ad2d8d27/test_to" # 输出文件夹
    os.makedirs(output_dir, exist_ok=True)
    for each_dir in os.listdir(input_dir):
        dir_path = os.path.join(input_dir, each_dir)
        split_series_time(dir_path, output_dir)   #标示最先拍摄的series
        split_series_thickness(dir_path, output_dir)    #标示每个series的厚度
